import numpy as np

import phe as paillier
import math

# 公钥pk,私钥sk
pk, sk = paillier.generate_paillier_keypair()

# 最大明文长度
max_bit = 1024 # math.floor(math.log(pk.n, 2))

# 精度prec,填充pad
prec, pad = 8, 0   #31, 1

# 单个数据大小
data_size = prec + pad

# 明文可容纳的数据数目
t = math.floor(max_bit / data_size) -1

# 浮点数输入的上下界
LOG2_BOUND = 0  # 小数点？
bound = 2 ** LOG2_BOUND


# 获取常数
def get_constant():
    return t, bound


# 打包加密
# arg1: 浮点数列表
def get_encrypted_bits(r_list):
    pack_bits = ''
    for i, r in enumerate(r_list):
        # float to int
        int_r = r #math.floor((r + bound) * 2 ** (prec - LOG2_BOUND - 1))  # params_grad = ((torch.cat(params_grad_list, 0) + bound) * 2 ** prec).long().cuda()
        # packing data in encryption
        # tmp = bin(int_r)[2:]
        # print(len(tmp))
        tmp = bin(int_r)[2:].zfill(prec)
        l = len(tmp)

        # if l >= data_size:
        #     print(r, l)
        pack_bits += tmp
        # print(len(pack_bits))
        if i != t - 1:
            if l <= prec and pad > 0:
                zero_pad = '0'
                pack_bits += zero_pad.zfill(pad)
    print(len(pack_bits))

    ct = pk.encrypt(int('0b' + pack_bits, 2))
    print(ct.ciphertext())
    print(len(bin(ct.ciphertext())[2:]))
    print(len(str(ct.ciphertext())))

    return pk.encrypt(int('0b' + pack_bits, 2))


# 解密拆包，
# arg1: 打包的密文
# arg2: 加数数目
def get_decrypted_number(pack_ciphertext, n):
    total_sum_bits = bin(sk.decrypt(pack_ciphertext))[2:].zfill(t * data_size)
    idx = 0
    sum_list = []
    for i in range(t):
        int_sum = int('0b' + total_sum_bits[idx: idx + data_size], 2)
        # sum_list.append(int_sum / 2 ** (prec - LOG2_BOUND - 1) - n * bound)  #int_r / 2 ** (prec - LOG2_BOUND - 1) - 1 * bound
        sum_list.append(int_sum)
        idx += data_size
    return sum_list


if __name__ == '__main__':
    import random
    t, bound = get_constant()

    image = "D:/FACE_TEST.png"
    import cv2
    imdata = cv2.imread(image)

    # from sklearn.datasets import fetch_lfw_people
    # lfw_people = fetch_lfw_people(min_faces_per_person=70, resize=0.4)
    # n_samples, h, w = lfw_people.images.shape
    # X = lfw_people.data
    # n_features = X.shape[1]
    # imdata = X[0]

    imdata=imdata.reshape(-1)

    batch_encrypted_list = []
    # values = [0.15, 0.1]
    i = 0
    while i < len(imdata):
        r_list = list(imdata[i: min(i+t, len(imdata))])
        # r_list = [0, 128, 255] #2**31
        # r_list = [int(x*1e8) for x in [0.15, 10.1]]
        for j in range(t-len(r_list)):
            # random float r
            r = 0 # random.uniform(-bound, bound)
            r_list.append(r)
        batch_encrypted1 = get_encrypted_bits(r_list)
        batch_encrypted_list.append(batch_encrypted1)
        batch_decrypted1 = get_decrypted_number(batch_encrypted1, 1)
        i += t

    # batch_encrypted1 = sum([batch_encrypted1])
    batch_decrypted1 = get_decrypted_number(batch_encrypted1, 1)
    print(batch_decrypted1)